{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 计算图分区"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "import tvm\n",
    "from tvm.relay.backend import te_compiler\n",
    "from tvm.relay.backend.runtime import Runtime\n",
    "import tvm.relay.testing\n",
    "from tvm import relay\n",
    "from tvm import runtime as tvm_runtime\n",
    "from tvm.relay import transform\n",
    "from tvm.contrib import utils\n",
    "from tvm.relay.expr_functor import ExprMutator\n",
    "from tvm.relay.op.annotation import compiler_begin, compiler_end\n",
    "from tvm.relay.op.contrib.register import get_pattern_table\n",
    "from tvm.relay.build_module import bind_params_by_name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "# Leverage the pass manager to write a simple allowed list based annotator\n",
    "@transform.function_pass(opt_level=0)\n",
    "class AllowedListAnnotator:\n",
    "    def __init__(self, op_list, compiler):\n",
    "        assert isinstance(op_list, (list, tuple, set))\n",
    "        self.op_list = op_list\n",
    "        self.compiler = compiler\n",
    "\n",
    "    def transform_function(self, func, mod, dev):\n",
    "\n",
    "        annotator = self\n",
    "\n",
    "        class Annotator(tvm.relay.ExprMutator):\n",
    "            def visit_call(self, call):\n",
    "                op_name = call.op.name\n",
    "                if op_name in annotator.op_list:\n",
    "                    new_args = []\n",
    "                    for arg in call.args:\n",
    "                        ann = compiler_begin(super().visit(arg), annotator.compiler)\n",
    "                        new_args.append(ann)\n",
    "                    new_call = relay.Call(call.op, new_args, call.attrs, call.type_args)\n",
    "                    return compiler_end(new_call, annotator.compiler)\n",
    "                else:\n",
    "                    return super().visit_call(call)\n",
    "\n",
    "        return Annotator().visit(func)\n",
    "\n",
    "\n",
    "class WholeGraphAnnotator(ExprMutator):\n",
    "    \"\"\"\n",
    "    An annotator that creates a compiler for an entire graph.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, compiler):\n",
    "        super(WholeGraphAnnotator, self).__init__()\n",
    "        self.compiler = compiler\n",
    "        self.last_call = True\n",
    "\n",
    "    def visit_call(self, call):\n",
    "        curr_last = self.last_call\n",
    "        self.last_call = False\n",
    "\n",
    "        params = []\n",
    "        for arg in call.args:\n",
    "            param = super().visit(arg)\n",
    "            if isinstance(param, relay.expr.Var):\n",
    "                param = compiler_begin(param, self.compiler)\n",
    "            params.append(param)\n",
    "\n",
    "        new_call = relay.Call(call.op, params, call.attrs)\n",
    "        if curr_last:\n",
    "            new_call = compiler_end(new_call, self.compiler)\n",
    "        return new_call\n",
    "\n",
    "\n",
    "class MobileNetAnnotator(ExprMutator):\n",
    "    \"\"\"\n",
    "    Annotate mobilenet until global_avg_pool.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, compiler):\n",
    "        super(MobileNetAnnotator, self).__init__()\n",
    "        self.compiler = compiler\n",
    "        self.compiler_open = False\n",
    "\n",
    "    def visit_call(self, call):\n",
    "\n",
    "        if call.op.name == \"nn.global_avg_pool2d\":\n",
    "            self.compiler_open = True\n",
    "        compiler_open = self.compiler_open\n",
    "\n",
    "        params = []\n",
    "        for arg in call.args:\n",
    "            param = super().visit(arg)\n",
    "            if call.op.name == \"nn.global_avg_pool2d\":\n",
    "                param = compiler_end(param, self.compiler)\n",
    "            if compiler_open and isinstance(param, relay.expr.Var):\n",
    "                param = compiler_begin(param, self.compiler)\n",
    "            params.append(param)\n",
    "\n",
    "        new_call = relay.Call(call.op, params, call.attrs)\n",
    "        return new_call"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "import os\n",
    "def check_result(\n",
    "    mod,\n",
    "    map_inputs,\n",
    "    out_shape,\n",
    "    result,\n",
    "    tol=1e-5,\n",
    "    target=\"llvm\",\n",
    "    device=tvm.cpu(),\n",
    "    params=None,\n",
    "    runtime=Runtime(\"cpp\"),\n",
    "):\n",
    "    def update_lib(lib):\n",
    "        test_dir = os.path.dirname(os.path.realpath(os.path.expanduser(\"__file__\")))\n",
    "        source_dir = os.path.join(test_dir, \"..\", \"..\", \"..\")\n",
    "        contrib_path = os.path.join(source_dir, \"src\", \"runtime\", \"contrib\")\n",
    "\n",
    "        kwargs = {}\n",
    "        kwargs[\"options\"] = [\"-O2\", \"-std=c++17\", \"-I\" + contrib_path]\n",
    "        tmp_path = utils.tempdir()\n",
    "        lib_name = \"lib.so\"\n",
    "        lib_path = tmp_path.relpath(lib_name)\n",
    "        lib.export_library(lib_path, fcompile=False, **kwargs)\n",
    "        lib = tvm_runtime.load_module(lib_path)\n",
    "\n",
    "        return lib\n",
    "\n",
    "    def check_vm_result():\n",
    "        te_compiler.get().clear()\n",
    "        with tvm.transform.PassContext(opt_level=3):\n",
    "            exe = relay.vm.compile(mod, target=target, params=params)\n",
    "        code, lib = exe.save()\n",
    "        lib = update_lib(lib)\n",
    "        exe = tvm_runtime.vm.Executable.load_exec(code, lib)\n",
    "        vm = tvm_runtime.vm.VirtualMachine(exe, device)\n",
    "        outs = vm.run(**map_inputs)\n",
    "        outs = outs if isinstance(outs, tvm_runtime.container.ADT) else [outs]\n",
    "        results = result if isinstance(result, list) else [result]\n",
    "        for out, ref in zip(outs, results):\n",
    "            tvm.testing.assert_allclose(out.numpy(), ref, rtol=tol, atol=tol)\n",
    "\n",
    "    def check_graph_executor_result():\n",
    "        te_compiler.get().clear()\n",
    "        with tvm.transform.PassContext(opt_level=3):\n",
    "            json, lib, param = relay.build(mod, target=target, params=params, runtime=runtime)\n",
    "        lib = update_lib(lib)\n",
    "        rt_mod = tvm.contrib.graph_executor.create(json, lib, device)\n",
    "\n",
    "        for name, data in map_inputs.items():\n",
    "            rt_mod.set_input(name, data)\n",
    "        rt_mod.set_input(**param)\n",
    "        rt_mod.run()\n",
    "\n",
    "        out_shapes = out_shape if isinstance(out_shape, list) else [out_shape]\n",
    "        results = result if isinstance(result, list) else [result]\n",
    "\n",
    "        for idx, shape in enumerate(out_shapes):\n",
    "            out = tvm.nd.empty(shape, device=device)\n",
    "            out = rt_mod.get_output(idx, out)\n",
    "            tvm.testing.assert_allclose(out.numpy(), results[idx], rtol=tol, atol=tol)\n",
    "\n",
    "    check_vm_result()\n",
    "    check_graph_executor_result()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 测试外部 C 编译器处理单一算子"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[17:36:37] /media/pc/data/lxw/ai/tvm/src/relay/backend/vm/compiler.cc:1199: All lowered functions have been build by BYOC -- generating an empty TVM module\n",
      "/tmp/ipykernel_2738903/3034817802.py:45: DeprecationWarning: legacy graph executor behavior of producing json / lib / params will be removed in the next release. Please see documents of tvm.contrib.graph_executor.GraphModule for the  new recommended usage.\n",
      "  json, lib, param = relay.build(mod, target=target, params=params, runtime=runtime)\n"
     ]
    }
   ],
   "source": [
    "@transform.function_pass(opt_level=0)\n",
    "class MyAnnotator:\n",
    "    def transform_function(self, func, mod, dev):\n",
    "        class Annotator(tvm.relay.ExprMutator):\n",
    "            def visit_call(self, call):\n",
    "                new_args = []\n",
    "                for arg in call.args:\n",
    "                    ann = compiler_begin(self.visit(arg), \"ccompiler\")\n",
    "                    new_args.append(ann)\n",
    "                new_call = relay.Call(call.op, new_args)\n",
    "                return compiler_end(new_call, \"ccompiler\")\n",
    "\n",
    "        return Annotator().visit(func)\n",
    "\n",
    "x = relay.var(\"x\", shape=(8, 8))\n",
    "y = relay.var(\"y\", shape=(8, 8))\n",
    "z = x + y\n",
    "f = relay.Function([x, y], z)\n",
    "x_data = np.random.rand(8, 8).astype(\"float32\")\n",
    "y_data = np.random.rand(8, 8).astype(\"float32\")\n",
    "mod = tvm.IRModule()\n",
    "mod[\"main\"] = f\n",
    "mod = MyAnnotator()(mod)\n",
    "mod = transform.PartitionGraph()(mod)\n",
    "\n",
    "check_result(mod, {\"x\": x_data, \"y\": y_data}, (8, 8), x_data + y_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 测试外部 C 编译器处理默认算子"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def set_func_attr(func, compile_name, symbol_name):\n",
    "    func = func.with_attr(\"Primitive\", tvm.tir.IntImm(\"int32\", 1))\n",
    "    func = func.with_attr(\"Inline\", tvm.tir.IntImm(\"int32\", 1))\n",
    "    func = func.with_attr(\"Compiler\", compile_name)\n",
    "    func = func.with_attr(\"global_symbol\", symbol_name)\n",
    "    return func"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_2738903/3034817802.py:45: DeprecationWarning: legacy graph executor behavior of producing json / lib / params will be removed in the next release. Please see documents of tvm.contrib.graph_executor.GraphModule for the  new recommended usage.\n",
      "  json, lib, param = relay.build(mod, target=target, params=params, runtime=runtime)\n"
     ]
    }
   ],
   "source": [
    "x = relay.var(\"x\", shape=(8, 8))\n",
    "y = relay.var(\"y\", shape=(8, 8))\n",
    "add = x + y\n",
    "log = relay.log(add)\n",
    "exp = relay.exp(add)\n",
    "concat = relay.concatenate([log, exp], axis=0)\n",
    "f = relay.Function([x, y], concat)\n",
    "mod = tvm.IRModule()\n",
    "mod[\"main\"] = f\n",
    "mod = AllowedListAnnotator([\"add\", \"subtract\", \"multiply\"], \"ccompiler\")(mod)\n",
    "mod = transform.PartitionGraph()(mod)\n",
    "fused_mod = transform.FuseOps(2)(mod)\n",
    "\n",
    "x_data = np.random.rand(8, 8).astype(\"float32\")\n",
    "y_data = np.random.rand(8, 8).astype(\"float32\")\n",
    "np_add = x_data + y_data\n",
    "res = np.concatenate([np.log(np_add), np.exp(np_add)])\n",
    "check_result(mod, {\"x\": x_data, \"y\": y_data}, (16, 8), res)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 测试外部编译器处理经过清理的算子"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@main</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>x: Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>y: Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">8</span>), float32] {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> <span style=\"color: #AA22FF\">@tvmgen_default_unsanitary_name___main_0</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>x, <span style=\"color: #AA22FF; font-weight: bold\">%</span>y) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">4</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> fn (<span style=\"color: #AA22FF; font-weight: bold\">%</span>p0: Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, Primitive<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">8</span>), float32] {\n",
       "    <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> log(<span style=\"color: #AA22FF; font-weight: bold\">%</span>p0) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "    <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> exp(<span style=\"color: #AA22FF; font-weight: bold\">%</span>p0) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "    <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> (<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>(Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32], Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32]) <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "    concatenate(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "  } <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>fn (Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32]) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">4</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "}\n",
       "\n",
       "<span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@tvmgen_default_unsanitary_name___main_0</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>unsanitary<span style=\"color: #AA22FF; font-weight: bold\">-</span>name<span style=\"color: #AA22FF; font-weight: bold\">++</span>_0_i0: Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>unsanitary<span style=\"color: #AA22FF; font-weight: bold\">-</span>name<span style=\"color: #AA22FF; font-weight: bold\">++</span>_0_i1: Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, Compiler<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;unsanitary-name++&quot;</span>, Primitive<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>, Inline<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>, global_symbol<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;tvmgen_default_unsanitary_name___main_0&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] {\n",
       "  add(<span style=\"color: #AA22FF; font-weight: bold\">%</span>unsanitary<span style=\"color: #AA22FF; font-weight: bold\">-</span>name<span style=\"color: #AA22FF; font-weight: bold\">++</span>_0_i0, <span style=\"color: #AA22FF; font-weight: bold\">%</span>unsanitary<span style=\"color: #AA22FF; font-weight: bold\">-</span>name<span style=\"color: #AA22FF; font-weight: bold\">++</span>_0_i1) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "x = relay.var(\"x\", shape=(8, 8))\n",
    "y = relay.var(\"y\", shape=(8, 8))\n",
    "add = x + y\n",
    "log = relay.log(add)\n",
    "exp = relay.exp(add)\n",
    "concat = relay.concatenate([log, exp], axis=0)\n",
    "f = relay.Function([x, y], concat)\n",
    "mod = tvm.IRModule()\n",
    "mod[\"main\"] = f\n",
    "mod = AllowedListAnnotator([\"add\", \"subtract\", \"multiply\"], \"unsanitary-name++\")(mod)\n",
    "mod = transform.PartitionGraph()(mod)\n",
    "fused_mod = transform.FuseOps(2)(mod)\n",
    "fused_mod.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 测试外部 C 编译器处理多个函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_2738903/3034817802.py:45: DeprecationWarning: legacy graph executor behavior of producing json / lib / params will be removed in the next release. Please see documents of tvm.contrib.graph_executor.GraphModule for the  new recommended usage.\n",
      "  json, lib, param = relay.build(mod, target=target, params=params, runtime=runtime)\n"
     ]
    }
   ],
   "source": [
    "x = relay.var(\"x\", shape=(8, 8))\n",
    "y = relay.var(\"y\", shape=(8, 8))\n",
    "add = x + y\n",
    "log = relay.log(add)\n",
    "exp = relay.exp(add)\n",
    "concat = relay.concatenate([log, exp], axis=0)\n",
    "f = relay.Function([x, y], concat)\n",
    "mod = tvm.IRModule()\n",
    "mod[\"main\"] = f\n",
    "# define second function\n",
    "a = relay.var(\"a\", shape=(16, 16))\n",
    "b = relay.var(\"b\", shape=(16, 16))\n",
    "add = a + b\n",
    "log = relay.log(add)\n",
    "exp = relay.exp(add)\n",
    "concat = relay.concatenate([log, exp], axis=0)\n",
    "f2 = relay.Function([a, b], concat)\n",
    "mod[\"subfunction\"] = f2\n",
    "mod = AllowedListAnnotator([\"add\", \"subtract\", \"multiply\"], \"ccompiler\")(mod)\n",
    "mod = transform.PartitionGraph()(mod)\n",
    "\n",
    "fused_mod = transform.FuseOps(2)(mod)\n",
    "\n",
    "x_data = np.random.rand(8, 8).astype(\"float32\")\n",
    "y_data = np.random.rand(8, 8).astype(\"float32\")\n",
    "np_add = x_data + y_data\n",
    "res = np.concatenate([np.log(np_add), np.exp(np_add)])\n",
    "check_result(mod, {\"x\": x_data, \"y\": y_data}, (16, 8), res)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 测试外部 C 编译器"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[17:36:40] /media/pc/data/lxw/ai/tvm/src/relay/backend/vm/compiler.cc:1199: All lowered functions have been build by BYOC -- generating an empty TVM module\n",
      "/tmp/ipykernel_2738903/3034817802.py:45: DeprecationWarning: legacy graph executor behavior of producing json / lib / params will be removed in the next release. Please see documents of tvm.contrib.graph_executor.GraphModule for the  new recommended usage.\n",
      "  json, lib, param = relay.build(mod, target=target, params=params, runtime=runtime)\n"
     ]
    }
   ],
   "source": [
    "x = relay.var(\"x\", shape=(2, 2))\n",
    "y = relay.var(\"y\", shape=(2, 2))\n",
    "z = x + x\n",
    "p = y * y\n",
    "f = relay.Function([x, y], p - z)\n",
    "x_data = np.random.rand(2, 2).astype(\"float32\")\n",
    "y_data = np.random.rand(2, 2).astype(\"float32\")\n",
    "mod = tvm.IRModule()\n",
    "mod[\"main\"] = f\n",
    "mod = AllowedListAnnotator([\"add\", \"subtract\", \"multiply\"], \"ccompiler\")(mod)\n",
    "mod = transform.PartitionGraph()(mod)\n",
    "\n",
    "check_result(mod, {\"x\": x_data, \"y\": y_data}, (2, 2), (y_data * y_data) - (x_data + x_data))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 测试函数提升"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "def partition():\n",
    "    data = relay.var(\"data\", relay.TensorType((1, 3, 224, 224), \"float32\"))\n",
    "    weight = relay.var(\"weight\", relay.TensorType((16, 3, 3, 3), \"float32\"))\n",
    "    bn_gamma = relay.var(\"bn_gamma\", relay.TensorType((16,), \"float32\"))\n",
    "    bn_beta = relay.var(\"bn_beta\", relay.TensorType((16,), \"float32\"))\n",
    "    bn_mmean = relay.var(\"bn_mean\", relay.TensorType((16,), \"float32\"))\n",
    "    bn_mvar = relay.var(\"bn_var\", relay.TensorType((16,), \"float32\"))\n",
    "\n",
    "    conv = relay.nn.conv2d(\n",
    "        data=data, weight=weight, kernel_size=(3, 3), channels=16, padding=(1, 1)\n",
    "    )\n",
    "    bn_output = relay.nn.batch_norm(conv, bn_gamma, bn_beta, bn_mmean, bn_mvar)\n",
    "\n",
    "    func = relay.Function(\n",
    "        [data, weight, bn_gamma, bn_beta, bn_mmean, bn_mvar], bn_output.astuple()\n",
    "    )\n",
    "    mod = tvm.IRModule()\n",
    "    mod[\"main\"] = func\n",
    "    mod = relay.transform.InferType()(mod)\n",
    "    op_list = [\"nn.batch_norm\", \"nn.conv2d\"]\n",
    "    mod = AllowedListAnnotator(op_list, \"test_compiler\")(mod)\n",
    "\n",
    "    opt_pass = tvm.transform.Sequential(\n",
    "        [\n",
    "            transform.InferType(),\n",
    "            transform.PartitionGraph(),\n",
    "            transform.SimplifyInference(),\n",
    "            transform.FoldConstant(),\n",
    "            transform.AlterOpLayout(),\n",
    "        ]\n",
    "    )\n",
    "\n",
    "    with tvm.transform.PassContext(opt_level=3):\n",
    "        mod = opt_pass(mod)\n",
    "\n",
    "    return mod"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@main</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>data: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">224</span>, <span style=\"color: #008000\">224</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">224</span>, <span style=\"color: #008000\">224</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>weight: Tensor[(<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>bn_gamma: Tensor[(<span style=\"color: #008000\">16</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">16</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>bn_beta: Tensor[(<span style=\"color: #008000\">16</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">16</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>bn_mean: Tensor[(<span style=\"color: #008000\">16</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">16</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>bn_var: Tensor[(<span style=\"color: #008000\">16</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">16</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> (Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">224</span>, <span style=\"color: #008000\">224</span>), float32], Tensor[(<span style=\"color: #008000\">16</span>), float32], Tensor[(<span style=\"color: #008000\">16</span>), float32]) {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> <span style=\"color: #AA22FF\">@tvmgen_default_test_compiler_main_0</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>data, <span style=\"color: #AA22FF; font-weight: bold\">%</span>weight) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">224</span>, <span style=\"color: #008000\">224</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF\">@tvmgen_default_test_compiler_main_2</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>bn_gamma, <span style=\"color: #AA22FF; font-weight: bold\">%</span>bn_beta, <span style=\"color: #AA22FF; font-weight: bold\">%</span>bn_mean, <span style=\"color: #AA22FF; font-weight: bold\">%</span>bn_var) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>(Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">224</span>, <span style=\"color: #008000\">224</span>), float32], Tensor[(<span style=\"color: #008000\">16</span>), float32], Tensor[(<span style=\"color: #008000\">16</span>), float32]) <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "}\n",
       "\n",
       "<span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@tvmgen_default_test_compiler_main_0</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>test_compiler_0_i0: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">224</span>, <span style=\"color: #008000\">224</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">224</span>, <span style=\"color: #008000\">224</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>test_compiler_0_i1: Tensor[(<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, Compiler<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;test_compiler&quot;</span>, Primitive<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>, Inline<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>, global_symbol<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;tvmgen_default_test_compiler_main_0&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">224</span>, <span style=\"color: #008000\">224</span>), float32] {\n",
       "  nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>conv2d(<span style=\"color: #AA22FF; font-weight: bold\">%</span>test_compiler_0_i0, <span style=\"color: #AA22FF; font-weight: bold\">%</span>test_compiler_0_i1, padding<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>], channels<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">16</span>, kernel_size<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>]) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">224</span>, <span style=\"color: #008000\">224</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "}\n",
       "\n",
       "<span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@tvmgen_default_test_compiler_main_2</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>test_compiler_2_i0: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">224</span>, <span style=\"color: #008000\">224</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">224</span>, <span style=\"color: #008000\">224</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>test_compiler_2_i1: Tensor[(<span style=\"color: #008000\">16</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">16</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>test_compiler_2_i2: Tensor[(<span style=\"color: #008000\">16</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">16</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>test_compiler_2_i3: Tensor[(<span style=\"color: #008000\">16</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">16</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>test_compiler_2_i4: Tensor[(<span style=\"color: #008000\">16</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">16</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, Compiler<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;test_compiler&quot;</span>, Primitive<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>, Inline<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>, global_symbol<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;tvmgen_default_test_compiler_main_2&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> (Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">224</span>, <span style=\"color: #008000\">224</span>), float32], Tensor[(<span style=\"color: #008000\">16</span>), float32], Tensor[(<span style=\"color: #008000\">16</span>), float32]) {\n",
       "  nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>batch_norm(<span style=\"color: #AA22FF; font-weight: bold\">%</span>test_compiler_2_i0, <span style=\"color: #AA22FF; font-weight: bold\">%</span>test_compiler_2_i1, <span style=\"color: #AA22FF; font-weight: bold\">%</span>test_compiler_2_i2, <span style=\"color: #AA22FF; font-weight: bold\">%</span>test_compiler_2_i3, <span style=\"color: #AA22FF; font-weight: bold\">%</span>test_compiler_2_i4) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>(Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">224</span>, <span style=\"color: #008000\">224</span>), float32], Tensor[(<span style=\"color: #008000\">16</span>), float32], Tensor[(<span style=\"color: #008000\">16</span>), float32]) <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "partitioned = partition()\n",
    "partitioned.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 测试函数提升内联"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@main</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>data: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">224</span>, <span style=\"color: #008000\">224</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">224</span>, <span style=\"color: #008000\">224</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>bn_gamma: Tensor[(<span style=\"color: #008000\">16</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">16</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>bn_beta: Tensor[(<span style=\"color: #008000\">16</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">16</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>bn_mean: Tensor[(<span style=\"color: #008000\">16</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">16</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>bn_var: Tensor[(<span style=\"color: #008000\">16</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">16</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> (Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">224</span>, <span style=\"color: #008000\">224</span>), float32], Tensor[(<span style=\"color: #008000\">16</span>), float32], Tensor[(<span style=\"color: #008000\">16</span>), float32]) {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> fn (<span style=\"color: #AA22FF; font-weight: bold\">%</span>test_compiler_0_i0: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">224</span>, <span style=\"color: #008000\">224</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">224</span>, <span style=\"color: #008000\">224</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>test_compiler_0_i1: Tensor[(<span style=\"color: #008000\">16</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">16</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>test_compiler_0_i2: Tensor[(<span style=\"color: #008000\">16</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">16</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>test_compiler_0_i3: Tensor[(<span style=\"color: #008000\">16</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">16</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>test_compiler_0_i4: Tensor[(<span style=\"color: #008000\">16</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">16</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, Compiler<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;test_compiler&quot;</span>, Primitive<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>, Inline<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>, global_symbol<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;tvmgen_default_test_compiler_main_0&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> (Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">224</span>, <span style=\"color: #008000\">224</span>), float32], Tensor[(<span style=\"color: #008000\">16</span>), float32], Tensor[(<span style=\"color: #008000\">16</span>), float32]) {\n",
       "    nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>batch_norm(<span style=\"color: #AA22FF; font-weight: bold\">%</span>test_compiler_0_i0, <span style=\"color: #AA22FF; font-weight: bold\">%</span>test_compiler_0_i1, <span style=\"color: #AA22FF; font-weight: bold\">%</span>test_compiler_0_i2, <span style=\"color: #AA22FF; font-weight: bold\">%</span>test_compiler_0_i3, <span style=\"color: #AA22FF; font-weight: bold\">%</span>test_compiler_0_i4) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>(Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">224</span>, <span style=\"color: #008000\">224</span>), float32], Tensor[(<span style=\"color: #008000\">16</span>), float32], Tensor[(<span style=\"color: #008000\">16</span>), float32]) <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "  };\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>data, <span style=\"color: #AA22FF; font-weight: bold\">%</span>bn_gamma, <span style=\"color: #AA22FF; font-weight: bold\">%</span>bn_beta, <span style=\"color: #AA22FF; font-weight: bold\">%</span>bn_mean, <span style=\"color: #AA22FF; font-weight: bold\">%</span>bn_var)\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def partition():\n",
    "    data = relay.var(\"data\", relay.TensorType((1, 16, 224, 224), \"float32\"))\n",
    "    bn_gamma = relay.var(\"bn_gamma\", relay.TensorType((16,), \"float32\"))\n",
    "    bn_beta = relay.var(\"bn_beta\", relay.TensorType((16,), \"float32\"))\n",
    "    bn_mmean = relay.var(\"bn_mean\", relay.TensorType((16,), \"float32\"))\n",
    "    bn_mvar = relay.var(\"bn_var\", relay.TensorType((16,), \"float32\"))\n",
    "\n",
    "    bn_output = relay.nn.batch_norm(data, bn_gamma, bn_beta, bn_mmean, bn_mvar)\n",
    "\n",
    "    func = relay.Function([data, bn_gamma, bn_beta, bn_mmean, bn_mvar], bn_output.astuple())\n",
    "    mod = tvm.IRModule()\n",
    "    mod[\"main\"] = func\n",
    "    op_list = [\"nn.batch_norm\", \"nn.conv2d\"]\n",
    "    mod = AllowedListAnnotator(op_list, \"test_compiler\")(mod)\n",
    "\n",
    "    opt_pass = tvm.transform.Sequential(\n",
    "        [\n",
    "            transform.InferType(),\n",
    "            transform.PartitionGraph(),\n",
    "            transform.SimplifyInference(),\n",
    "            transform.FoldConstant(),\n",
    "            transform.AlterOpLayout(),\n",
    "            transform.Inline(),\n",
    "        ]\n",
    "    )\n",
    "\n",
    "    with tvm.transform.PassContext(opt_level=3):\n",
    "        mod = opt_pass(mod)\n",
    "\n",
    "    return mod\n",
    "\n",
    "partitioned = partition()\n",
    "partitioned.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 测试常量传播"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_2738903/3034817802.py:45: DeprecationWarning: legacy graph executor behavior of producing json / lib / params will be removed in the next release. Please see documents of tvm.contrib.graph_executor.GraphModule for the  new recommended usage.\n",
      "  json, lib, param = relay.build(mod, target=target, params=params, runtime=runtime)\n"
     ]
    }
   ],
   "source": [
    "ones = np.ones(shape=(8, 8), dtype=\"float32\")\n",
    "x = relay.var(\"x\", shape=(8, 8))\n",
    "y = relay.var(\"y\", shape=(8, 8))\n",
    "add = x + y\n",
    "log = relay.log(add)\n",
    "f = relay.Function([x, y], log)\n",
    "f = bind_params_by_name(f, {\"x\": tvm.nd.array(ones)})\n",
    "mod = tvm.IRModule()\n",
    "mod[\"main\"] = f\n",
    "mod = AllowedListAnnotator([\"add\"], \"ccompiler\")(mod)\n",
    "mod = transform.PartitionGraph()(mod)\n",
    "mod = relay.transform.InferType()(mod)\n",
    "\n",
    "y_data = np.random.rand(8, 8).astype(\"float32\")\n",
    "np_add = ones + y_data\n",
    "check_result(mod, {\"y\": y_data}, (8, 8), np.log(np_add))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 测试多输出"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "def create_graph():\n",
    "    data = relay.var(\"data\", relay.TensorType((1, 3, 224, 224), \"float32\"))\n",
    "    weight = relay.var(\"weight\", relay.TensorType((16, 3, 3, 3), \"float32\"))\n",
    "    bn_gamma = relay.var(\"bn_gamma\", relay.TensorType((16,), \"float32\"))\n",
    "    bn_beta = relay.var(\"bn_beta\", relay.TensorType((16,), \"float32\"))\n",
    "    bn_mean = relay.var(\"bn_mean\", relay.TensorType((16,), \"float32\"))\n",
    "    bn_var = relay.var(\"bn_var\", relay.TensorType((16,), \"float32\"))\n",
    "\n",
    "    data_cb = compiler_begin(data, \"test_target\")\n",
    "    weight_cb = compiler_begin(weight, \"test_target\")\n",
    "    bn_gamma_cb = compiler_begin(bn_gamma, \"test_target\")\n",
    "    bn_beta_cb = compiler_begin(bn_beta, \"test_target\")\n",
    "    bn_mean_cb = compiler_begin(bn_mean, \"test_target\")\n",
    "    bn_var_cb = compiler_begin(bn_var, \"test_target\")\n",
    "\n",
    "    conv_o = relay.nn.conv2d(\n",
    "        data=data_cb, weight=weight_cb, kernel_size=(3, 3), channels=16, padding=(1, 1)\n",
    "    )\n",
    "\n",
    "    bn_o = relay.nn.batch_norm(conv_o, bn_gamma_cb, bn_beta_cb, bn_mean_cb, bn_var_cb)\n",
    "\n",
    "    relu_o = relay.nn.relu(bn_o[0])\n",
    "    relu_o_ce = compiler_end(relu_o, \"test_target\")\n",
    "\n",
    "    bn_omean = bn_o[1]\n",
    "    rebn_omean_ce = compiler_end(bn_omean, \"test_target\")\n",
    "    bn_ovar = bn_o[2]\n",
    "    bn_ovar_ce = compiler_end(bn_ovar, \"test_target\")\n",
    "\n",
    "    dummy_mean_abs = relay.abs(rebn_omean_ce)\n",
    "    dummy_ovar_abs = relay.abs(bn_ovar_ce)\n",
    "    dummy_tuple = relay.Tuple((relu_o_ce, dummy_mean_abs, dummy_ovar_abs))\n",
    "\n",
    "    func = relay.Function([data, weight, bn_gamma, bn_beta, bn_mean, bn_var], dummy_tuple)\n",
    "    return func\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@main</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>data: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">224</span>, <span style=\"color: #008000\">224</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">224</span>, <span style=\"color: #008000\">224</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>weight: Tensor[(<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>bn_gamma: Tensor[(<span style=\"color: #008000\">16</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">16</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>bn_beta: Tensor[(<span style=\"color: #008000\">16</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">16</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>bn_mean: Tensor[(<span style=\"color: #008000\">16</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">16</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>bn_var: Tensor[(<span style=\"color: #008000\">16</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">16</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> (Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">224</span>, <span style=\"color: #008000\">224</span>), float32], Tensor[(<span style=\"color: #008000\">16</span>), float32], Tensor[(<span style=\"color: #008000\">16</span>), float32]) {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> <span style=\"color: #AA22FF\">@tvmgen_default_test_target_main_0</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>data, <span style=\"color: #AA22FF; font-weight: bold\">%</span>weight, <span style=\"color: #AA22FF; font-weight: bold\">%</span>bn_gamma, <span style=\"color: #AA22FF; font-weight: bold\">%</span>bn_beta, <span style=\"color: #AA22FF; font-weight: bold\">%</span>bn_mean, <span style=\"color: #AA22FF; font-weight: bold\">%</span>bn_var) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>(Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">224</span>, <span style=\"color: #008000\">224</span>), float32], Tensor[(<span style=\"color: #008000\">16</span>), float32], Tensor[(<span style=\"color: #008000\">16</span>), float32]) <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0.1</span> <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">16</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0.2</span> <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">16</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0.0</span> <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">224</span>, <span style=\"color: #008000\">224</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">4</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> abs(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">16</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">5</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> abs(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">16</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  (<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">4</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">5</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>(Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">224</span>, <span style=\"color: #008000\">224</span>), float32], Tensor[(<span style=\"color: #008000\">16</span>), float32], Tensor[(<span style=\"color: #008000\">16</span>), float32]) <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "}\n",
       "\n",
       "<span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@tvmgen_default_test_target_main_0</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>test_target_0_i0: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">224</span>, <span style=\"color: #008000\">224</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">224</span>, <span style=\"color: #008000\">224</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>test_target_0_i1: Tensor[(<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>test_target_0_i2: Tensor[(<span style=\"color: #008000\">16</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">16</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>test_target_0_i3: Tensor[(<span style=\"color: #008000\">16</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">16</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>test_target_0_i4: Tensor[(<span style=\"color: #008000\">16</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">16</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>test_target_0_i5: Tensor[(<span style=\"color: #008000\">16</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">16</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, Compiler<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;test_target&quot;</span>, Primitive<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>, Inline<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>, global_symbol<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;tvmgen_default_test_target_main_0&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> (Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">224</span>, <span style=\"color: #008000\">224</span>), float32], Tensor[(<span style=\"color: #008000\">16</span>), float32], Tensor[(<span style=\"color: #008000\">16</span>), float32]) {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">6</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>conv2d(<span style=\"color: #AA22FF; font-weight: bold\">%</span>test_target_0_i0, <span style=\"color: #AA22FF; font-weight: bold\">%</span>test_target_0_i1, padding<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>], channels<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">16</span>, kernel_size<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>]) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">224</span>, <span style=\"color: #008000\">224</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">7</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>batch_norm(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">6</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>test_target_0_i2, <span style=\"color: #AA22FF; font-weight: bold\">%</span>test_target_0_i3, <span style=\"color: #AA22FF; font-weight: bold\">%</span>test_target_0_i4, <span style=\"color: #AA22FF; font-weight: bold\">%</span>test_target_0_i5) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>(Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">224</span>, <span style=\"color: #008000\">224</span>), float32], Tensor[(<span style=\"color: #008000\">16</span>), float32], Tensor[(<span style=\"color: #008000\">16</span>), float32]) <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">8</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">7.0</span> <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">224</span>, <span style=\"color: #008000\">224</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">9</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">8</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">224</span>, <span style=\"color: #008000\">224</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">10</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">7.1</span> <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">16</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">11</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">7.2</span> <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">16</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  (<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">9</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">10</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">11</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>(Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">224</span>, <span style=\"color: #008000\">224</span>), float32], Tensor[(<span style=\"color: #008000\">16</span>), float32], Tensor[(<span style=\"color: #008000\">16</span>), float32]) <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "mod = tvm.IRModule()\n",
    "mod[\"main\"] = create_graph()\n",
    "partitioned = transform.PartitionGraph()(mod)\n",
    "partitioned.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 测试混合单输出和多输出"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@main</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>data: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> <span style=\"color: #AA22FF\">@tvmgen_default_test_target_main_0</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>data) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>(Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32], Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32]) <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0.1</span> <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0.0</span> <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> tanh(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF\">@tvmgen_default_test_target_main_1</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "}\n",
       "\n",
       "<span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@tvmgen_default_test_target_main_0</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>test_target_0_i0: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, Compiler<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;test_target&quot;</span>, Primitive<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>, Inline<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>, global_symbol<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;tvmgen_default_test_target_main_0&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> (Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32], Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32]) {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">4</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> abs(<span style=\"color: #AA22FF; font-weight: bold\">%</span>test_target_0_i0) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">5</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">4</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  (<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">5</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">4</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>(Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32], Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32]) <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "}\n",
       "\n",
       "<span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@tvmgen_default_test_target_main_1</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>test_target_1_i0: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>test_target_1_i1: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, Compiler<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;test_target&quot;</span>, Primitive<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>, Inline<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>, global_symbol<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;tvmgen_default_test_target_main_1&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] {\n",
       "  add(<span style=\"color: #AA22FF; font-weight: bold\">%</span>test_target_1_i0, <span style=\"color: #AA22FF; font-weight: bold\">%</span>test_target_1_i1) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def create_graph():\n",
    "    data = relay.var(\"data\", shape=(10, 10))\n",
    "\n",
    "    cb_1 = compiler_begin(data, \"test_target\")\n",
    "    O_1 = relay.abs(cb_1)\n",
    "    ce_2 = compiler_end(O_1, \"test_target\")\n",
    "    O_2 = relay.nn.relu(O_1)\n",
    "    ce_3 = compiler_end(O_2, \"test_target\")\n",
    "\n",
    "    X = relay.tanh(ce_2)\n",
    "\n",
    "    cb_3 = compiler_begin(ce_3, \"test_target\")\n",
    "    cb_4 = compiler_begin(X, \"test_target\")\n",
    "    O_3 = relay.add(cb_3, cb_4)\n",
    "    ce_4 = compiler_end(O_3, \"test_target\")\n",
    "\n",
    "    func = relay.Function([data], ce_4)\n",
    "    return func\n",
    "\n",
    "mod = tvm.IRModule()\n",
    "mod[\"main\"] = create_graph()\n",
    "mod = transform.InferType()(mod)\n",
    "\n",
    "partitioned = transform.PartitionGraph()(mod)\n",
    "partitioned.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 测试一个输出被多次使用的情况"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "def expected_same_output_region():\n",
    "    mod = tvm.IRModule()\n",
    "    x = relay.var(\"x\", shape=(8, 8))\n",
    "    y = relay.var(\"y\", shape=(8, 8))\n",
    "    z = relay.var(\"z\", shape=(8, 8))\n",
    "    x0 = relay.var(\"x0\", shape=(8, 8))\n",
    "    y0 = relay.var(\"y0\", shape=(8, 8))\n",
    "    log = relay.log(x0)\n",
    "    sub = x0 - y0\n",
    "    mul = log * sub\n",
    "    # The partitioned graph contains log, subtract, and multiply\n",
    "    func = relay.Function([x0, y0], mul)\n",
    "    func = set_func_attr(func, \"ccompiler\", \"tvmgen_default_ccompiler_main_0\")\n",
    "    glb_0 = relay.GlobalVar(\"tvmgen_default_ccompiler_main_0\")\n",
    "    mod[glb_0] = func\n",
    "    mod = transform.InferType()(mod)\n",
    "\n",
    "    add = x + y\n",
    "    call = relay.Call(glb_0, [add, z])\n",
    "    main = relay.Function([x, y, z], call)\n",
    "    mod[\"main\"] = main\n",
    "    mod = transform.InferType()(mod)\n",
    "    return mod\n",
    "\n",
    "def expected_different_output_region():\n",
    "    mod = tvm.IRModule()\n",
    "    x = relay.var(\"x\", shape=(8, 8))\n",
    "    y = relay.var(\"y\", shape=(8, 8))\n",
    "    z = relay.var(\"z\", shape=(8, 8))\n",
    "\n",
    "    # The partitioned graph contains log\n",
    "    i0 = relay.var(\"i0\", shape=(8, 8))\n",
    "    log = relay.log(i0)\n",
    "    func = relay.Function([i0], log)\n",
    "    func = set_func_attr(func, \"ccompiler\", \"tvmgen_default_ccompiler_main_0\")\n",
    "    glb_0 = relay.GlobalVar(\"tvmgen_default_ccompiler_main_0\")\n",
    "    mod[glb_0] = func\n",
    "    mod = transform.InferType()(mod)\n",
    "\n",
    "    # The partitioned graph contains subtract\n",
    "    x0 = relay.var(\"x0\", shape=(8, 8))\n",
    "    y0 = relay.var(\"y0\", shape=(8, 8))\n",
    "    sub = x0 - y0\n",
    "    func = relay.Function([x0, y0], sub)\n",
    "    func = set_func_attr(func, \"ccompiler\", \"tvmgen_default_ccompiler_main_1\")\n",
    "    glb_1 = relay.GlobalVar(\"tvmgen_default_ccompiler_main_1\")\n",
    "    mod[glb_1] = func\n",
    "    mod = transform.InferType()(mod)\n",
    "\n",
    "    add = x + y\n",
    "    call_log = relay.Call(glb_0, [add])\n",
    "    call_sub = relay.Call(glb_1, [add, z])\n",
    "    main = relay.Function([x, y, z], call_log * call_sub)\n",
    "    mod[\"main\"] = main\n",
    "    mod = transform.InferType()(mod)\n",
    "    return mod\n",
    "\n",
    "def get_mod():\n",
    "    x = relay.var(\"x\", shape=(8, 8))\n",
    "    y = relay.var(\"y\", shape=(8, 8))\n",
    "    z = relay.var(\"z\", shape=(8, 8))\n",
    "    add = x + y\n",
    "    sub = add - z\n",
    "    log = relay.log(add)\n",
    "    sub1 = log * sub\n",
    "    f = relay.Function([x, y, z], sub1)\n",
    "    mod = tvm.IRModule()\n",
    "    mod[\"main\"] = f\n",
    "    return mod\n",
    "\n",
    "def test_same_output_region():\n",
    "    mod = get_mod()\n",
    "    mod = AllowedListAnnotator([\"subtract\", \"log\", \"multiply\"], \"ccompiler\")(mod)\n",
    "    mod = transform.MergeCompilerRegions()(mod)\n",
    "    mod = transform.PartitionGraph()(mod)\n",
    "\n",
    "    expected_mod = expected_same_output_region()\n",
    "    tvm.ir.assert_structural_equal(mod, expected_mod, map_free_vars=True)\n",
    "\n",
    "def test_different_output_region():\n",
    "    mod = get_mod()\n",
    "    mod = AllowedListAnnotator([\"subtract\", \"log\"], \"ccompiler\")(mod)\n",
    "    mod = transform.MergeCompilerRegions()(mod)\n",
    "    mod = transform.PartitionGraph()(mod)\n",
    "\n",
    "    expected_mod = expected_different_output_region()\n",
    "    tvm.ir.assert_structural_equal(mod, expected_mod, map_free_vars=True)\n",
    "\n",
    "test_same_output_region()\n",
    "test_different_output_region()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 测试重复输出的情况"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@main</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>data: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> (Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32], Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32], Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32]) {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> <span style=\"color: #AA22FF\">@tvmgen_default_test_duplicate_outputs_main_0</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>data) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> tanh(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> log(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  (<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>(Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32], Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32], Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32]) <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "}\n",
       "\n",
       "<span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@tvmgen_default_test_duplicate_outputs_main_0</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>test_duplicate_outputs_0_i0: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, Compiler<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;test_duplicate_outputs&quot;</span>, Primitive<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>, Inline<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>, global_symbol<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;tvmgen_default_test_duplicate_outputs_main_0&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] {\n",
       "  abs(<span style=\"color: #AA22FF; font-weight: bold\">%</span>test_duplicate_outputs_0_i0) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "target = \"test_duplicate_outputs\"\n",
    "\n",
    "@tvm.ir.register_op_attr(\"abs\", \"target.\" + target)\n",
    "def abs(expr):  # pylint: disable=unused-variable\n",
    "    return True\n",
    "\n",
    "def create_graph():\n",
    "    data = relay.var(\"data\", shape=(10, 10))\n",
    "    x = relay.abs(data)\n",
    "    out_1 = relay.nn.relu(x)\n",
    "    out_2 = relay.tanh(x)\n",
    "    out_3 = relay.log(x)\n",
    "    out = relay.Tuple([out_1, out_2, out_3])\n",
    "    func = relay.Function([data], out)\n",
    "    return func\n",
    "\n",
    "mod = tvm.IRModule()\n",
    "mod[\"main\"] = create_graph()\n",
    "\n",
    "seq = tvm.transform.Sequential(\n",
    "    [\n",
    "        transform.AnnotateTarget(target),\n",
    "        transform.MergeCompilerRegions(),\n",
    "        transform.PartitionGraph(),\n",
    "    ]\n",
    ")\n",
    "partitioned = seq(mod)\n",
    "partitioned.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 测试重复合并和 `TupleGetItem` 的情况"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@main</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>data: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>bn_gamma: Tensor[(<span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>bn_beta: Tensor[(<span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>bn_mean: Tensor[(<span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>bn_var: Tensor[(<span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> (Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32], Tensor[(<span style=\"color: #008000\">10</span>), float32], Tensor[(<span style=\"color: #008000\">10</span>), float32]) {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> <span style=\"color: #AA22FF\">@tvmgen_default_test_duplicate_merge_and_tuplegetitem_main_0</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>data, <span style=\"color: #AA22FF; font-weight: bold\">%</span>bn_gamma, <span style=\"color: #AA22FF; font-weight: bold\">%</span>bn_beta, <span style=\"color: #AA22FF; font-weight: bold\">%</span>bn_mean, <span style=\"color: #AA22FF; font-weight: bold\">%</span>bn_var) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>(Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32], Tensor[(<span style=\"color: #008000\">10</span>), float32]) <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0.1</span> <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0.0</span> <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> tanh(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">4</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> log(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  (<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">4</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>(Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32], Tensor[(<span style=\"color: #008000\">10</span>), float32], Tensor[(<span style=\"color: #008000\">10</span>), float32]) <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "}\n",
       "\n",
       "<span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@tvmgen_default_test_duplicate_merge_and_tuplegetitem_main_0</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>test_duplicate_merge_and_tuplegetitem_0_i0: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>test_duplicate_merge_and_tuplegetitem_0_i1: Tensor[(<span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>test_duplicate_merge_and_tuplegetitem_0_i2: Tensor[(<span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>test_duplicate_merge_and_tuplegetitem_0_i3: Tensor[(<span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>test_duplicate_merge_and_tuplegetitem_0_i4: Tensor[(<span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, Compiler<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;test_duplicate_merge_and_tuplegetitem&quot;</span>, Primitive<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>, Inline<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>, global_symbol<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;tvmgen_default_test_duplicate_merge_and_tuplegetitem_main_0&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> (Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32], Tensor[(<span style=\"color: #008000\">10</span>), float32]) {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">5</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>batch_norm(<span style=\"color: #AA22FF; font-weight: bold\">%</span>test_duplicate_merge_and_tuplegetitem_0_i0, <span style=\"color: #AA22FF; font-weight: bold\">%</span>test_duplicate_merge_and_tuplegetitem_0_i1, <span style=\"color: #AA22FF; font-weight: bold\">%</span>test_duplicate_merge_and_tuplegetitem_0_i2, <span style=\"color: #AA22FF; font-weight: bold\">%</span>test_duplicate_merge_and_tuplegetitem_0_i3, <span style=\"color: #AA22FF; font-weight: bold\">%</span>test_duplicate_merge_and_tuplegetitem_0_i4) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>(Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32], Tensor[(<span style=\"color: #008000\">10</span>), float32], Tensor[(<span style=\"color: #008000\">10</span>), float32]) <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">6</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">5.0</span> <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">7</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">6</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">8</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">5.1</span> <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  (<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">7</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">8</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>(Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32], Tensor[(<span style=\"color: #008000\">10</span>), float32]) <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "target = \"test_duplicate_merge_and_tuplegetitem\"\n",
    "\n",
    "@tvm.ir.register_op_attr(\"nn.batch_norm\", \"target.\" + target)\n",
    "def batch_norm(expr):  # pylint: disable=unused-variable\n",
    "    return True\n",
    "\n",
    "@tvm.ir.register_op_attr(\"nn.relu\", \"target.\" + target)\n",
    "def relu(expr):  # pylint: disable=unused-variable\n",
    "    return True\n",
    "\n",
    "def create_graph():\n",
    "    data = relay.var(\"data\", shape=(10, 10))\n",
    "    bn_gamma = relay.var(\"bn_gamma\")\n",
    "    bn_beta = relay.var(\"bn_beta\")\n",
    "    bn_mmean = relay.var(\"bn_mean\")\n",
    "    bn_mvar = relay.var(\"bn_var\")\n",
    "    x = relay.nn.batch_norm(data, bn_gamma, bn_beta, bn_mmean, bn_mvar)\n",
    "    out_1 = relay.nn.relu(x[0])\n",
    "    bn_out_1 = x[1]\n",
    "    out_2 = relay.tanh(bn_out_1)\n",
    "    out_3 = relay.log(bn_out_1)\n",
    "    out = relay.Tuple([out_1, out_2, out_3])\n",
    "    func = relay.Function([data, bn_gamma, bn_beta, bn_mmean, bn_mvar], out)\n",
    "    return func\n",
    "\n",
    "mod = tvm.IRModule()\n",
    "mod[\"main\"] = create_graph()\n",
    "mod = transform.InferType()(mod)\n",
    "\n",
    "seq = tvm.transform.Sequential(\n",
    "    [\n",
    "        transform.AnnotateTarget(target),\n",
    "        transform.MergeCompilerRegions(),\n",
    "        transform.PartitionGraph(),\n",
    "    ]\n",
    ")\n",
    "\n",
    "partitioned = seq(mod)\n",
    "partitioned.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 测试常量元组"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "@tvm.ir.register_op_attr(\"qnn.concatenate\", \"target.const_tuples\")\n",
    "def add(expr):  # pylint: disable=unused-variable\n",
    "    return True\n",
    "\n",
    "def create_graph():\n",
    "    a = relay.var(\"a\", shape=(10, 10), dtype=\"uint8\")\n",
    "    b = relay.var(\"b\", shape=(10, 10), dtype=\"uint8\")\n",
    "    a1 = relay.abs(a)\n",
    "\n",
    "    zeroi = relay.const(1, \"int32\")\n",
    "    zerof = relay.const(0, \"float32\")\n",
    "    con = relay.qnn.op.concatenate(\n",
    "        (a1, b),\n",
    "        input_scales=(zerof, zerof),\n",
    "        input_zero_points=(zeroi, zeroi),\n",
    "        output_scale=zerof,\n",
    "        output_zero_point=zeroi,\n",
    "        axis=1,\n",
    "    )\n",
    "\n",
    "    f = relay.Function([a, b], con)\n",
    "    mod = tvm.IRModule.from_expr(f)\n",
    "    mod = transform.InferType()(mod)\n",
    "    return mod\n",
    "\n",
    "seq = tvm.transform.Sequential(\n",
    "    [\n",
    "        transform.AnnotateTarget(\"const_tuples\"),\n",
    "        transform.InferType(),\n",
    "        transform.MergeCompilerRegions(),\n",
    "        transform.PartitionGraph(),\n",
    "    ]\n",
    ")\n",
    "\n",
    "partitioned = seq(create_graph())\n",
    "\n",
    "concat = partitioned[\"tvmgen_default_const_tuples_main_0\"].body\n",
    "assert type(concat.args[1]) == relay.Tuple\n",
    "assert type(concat.args[2]) == relay.Tuple\n",
    "assert type(concat.args[3]) == relay.Constant\n",
    "assert type(concat.args[4]) == relay.Constant\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 测试扁平化元组输出"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@main</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>a: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), uint8] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), uint8] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> (Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), uint8], Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">10</span>), uint8]) {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> <span style=\"color: #AA22FF\">@tvmgen_default_test_flatten_tuple_output_main_0</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>a) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>(Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">10</span>), uint8], Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">10</span>), uint8], Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">10</span>), uint8]) <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0.0</span> <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">10</span>), uint8] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0.1</span> <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">10</span>), uint8] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> (<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>(Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">10</span>), uint8], Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">10</span>), uint8]) <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">4</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0.2</span> <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">10</span>), uint8] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">5</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> concatenate(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), uint8] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">6</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">4</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">10</span>), uint8] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  (<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">5</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">6</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>(Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), uint8], Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">10</span>), uint8]) <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "}\n",
       "\n",
       "<span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@tvmgen_default_test_flatten_tuple_output_main_0</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>test_flatten_tuple_output_0_i0: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), uint8] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), uint8] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, Compiler<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;test_flatten_tuple_output&quot;</span>, Primitive<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>, Inline<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>, global_symbol<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;tvmgen_default_test_flatten_tuple_output_main_0&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> (Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">10</span>), uint8], Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">10</span>), uint8], Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">10</span>), uint8]) {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">7</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> split(<span style=\"color: #AA22FF; font-weight: bold\">%</span>test_flatten_tuple_output_0_i0, indices_or_sections<span style=\"color: #AA22FF; font-weight: bold\">=</span>meta[runtime<span style=\"color: #AA22FF; font-weight: bold\">.</span>BoxInt][<span style=\"color: #008000\">0</span>]) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>(Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">10</span>), uint8], Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">10</span>), uint8]) <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">8</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">7.0</span> <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">10</span>), uint8] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">9</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">7.0</span> <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">10</span>), uint8] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">10</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">7.1</span> <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">10</span>), uint8] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">11</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> abs(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">8</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">10</span>), uint8] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  (<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">9</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">10</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">11</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>(Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">10</span>), uint8], Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">10</span>), uint8], Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">10</span>), uint8]) <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "target = \"test_flatten_tuple_output\"\n",
    "\n",
    "@tvm.ir.register_op_attr(\"split\", \"target.\" + target)\n",
    "def split(expr):  # pylint: disable=unused-variable\n",
    "    return True\n",
    "\n",
    "@tvm.ir.register_op_attr(\"abs\", \"target.\" + target)\n",
    "def abs(expr):  # pylint: disable=unused-variable\n",
    "    return True\n",
    "\n",
    "def create_graph():\n",
    "    a = relay.var(\"a\", shape=(10, 10), dtype=\"uint8\")\n",
    "\n",
    "    a_split = relay.split(a, 2)\n",
    "    a_split_0 = relay.TupleGetItem(a_split.astuple(), 0)\n",
    "    a_split_0_abs = relay.abs(a_split_0)\n",
    "\n",
    "    a_con = relay.concatenate(a_split, 0)\n",
    "    a_split_0_relu = relay.nn.relu(a_split_0_abs)\n",
    "\n",
    "    out = relay.Tuple((a_con, a_split_0_relu))\n",
    "    f = relay.Function([a], out)\n",
    "    mod = tvm.IRModule.from_expr(f)\n",
    "    mod = transform.InferType()(mod)\n",
    "    return mod\n",
    "seq = tvm.transform.Sequential(\n",
    "    [\n",
    "        transform.AnnotateTarget(target),\n",
    "        transform.MergeCompilerRegions(),\n",
    "        transform.PartitionGraph(),\n",
    "    ]\n",
    ")\n",
    "\n",
    "partitioned = seq(create_graph())\n",
    "partitioned = transform.InferType()(partitioned)\n",
    "partitioned.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 测试元组输出的执行\n",
    "\n",
    "测试具有元组输出的子图的 C 代码生成和运行时。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[17:40:17] /media/pc/data/lxw/ai/tvm/src/relay/backend/vm/compiler.cc:1199: All lowered functions have been build by BYOC -- generating an empty TVM module\n",
      "/tmp/ipykernel_2738903/3034817802.py:45: DeprecationWarning: legacy graph executor behavior of producing json / lib / params will be removed in the next release. Please see documents of tvm.contrib.graph_executor.GraphModule for the  new recommended usage.\n",
      "  json, lib, param = relay.build(mod, target=target, params=params, runtime=runtime)\n"
     ]
    }
   ],
   "source": [
    "a = relay.var(\"a\", shape=(10, 10), dtype=\"float32\")\n",
    "b = relay.var(\"b\", shape=(10, 10), dtype=\"float32\")\n",
    "ba = relay.annotation.compiler_begin(a, \"ccompiler\")\n",
    "bb = relay.annotation.compiler_begin(b, \"ccompiler\")\n",
    "add = relay.add(ba, bb)\n",
    "sub = relay.subtract(ba, bb)\n",
    "out = relay.Tuple((add, sub))\n",
    "eout = relay.annotation.compiler_end(out, \"ccompiler\")\n",
    "func = relay.Function([a, b], eout)\n",
    "\n",
    "mod = tvm.IRModule()\n",
    "mod[\"main\"] = func\n",
    "mod = transform.InferType()(mod)\n",
    "mod = transform.PartitionGraph()(mod)\n",
    "\n",
    "a_data = np.random.rand(10, 10).astype(\"float32\")\n",
    "b_data = np.random.rand(10, 10).astype(\"float32\")\n",
    "\n",
    "check_result(\n",
    "    mod,\n",
    "    {\"a\": a_data, \"b\": b_data},\n",
    "    [(10, 10), (10, 10)],\n",
    "    [(a_data + b_data), (a_data - b_data)],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 测试外部优化"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "def Optimize(mod):\n",
    "    return relay.transform.FoldConstant()(mod)\n",
    "\n",
    "tvm.register_func(\"relay.ext.test_target.optimize\", Optimize)\n",
    "\n",
    "x = relay.var(\"x\", shape=(2, 2))\n",
    "y0 = relay.var(\"y0\", shape=(2, 2))\n",
    "y1 = relay.var(\"y1\", shape=(2, 2))\n",
    "yy0 = relay.annotation.compiler_begin(y0, \"test_target\")\n",
    "yy1 = relay.annotation.compiler_begin(y1, \"test_target\")\n",
    "z = yy0 + yy1\n",
    "end = relay.annotation.compiler_end(z, \"test_target\")\n",
    "f = relay.Function([x, y0, y1], end * x)\n",
    "c = np.ones(shape=(2, 2), dtype=\"float32\")\n",
    "f = bind_params_by_name(f, {\"y0\": tvm.nd.array(c), \"y1\": tvm.nd.array(c)})\n",
    "mod = tvm.IRModule()\n",
    "mod[\"main\"] = f\n",
    "mod = transform.InferType()(mod)\n",
    "mod = transform.PartitionGraph()(mod)\n",
    "\n",
    "try:\n",
    "    t0 = mod[\"tvmgen_default_test_target_main_0\"]\n",
    "except:\n",
    "    raise KeyError(\"test_target_main_0 not found\")\n",
    "\n",
    "assert isinstance(t0.body, relay.Constant)\n",
    "expected = np.empty([2, 2])\n",
    "expected.fill(2)\n",
    "tvm.testing.assert_allclose(t0.body.data.numpy(), expected, rtol=1e-5, atol=1e-5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 测试保留类型导入\n",
    "\n",
    "测试确保在 BYOC 管道中保留类型定义和导入。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tvm.relay.prelude import Prelude, StaticTensorArrayOps\n",
    "\n",
    "def run(dtype, shape):\n",
    "    mod = tvm.IRModule()\n",
    "    p = Prelude(mod)\n",
    "    static_tensor_array_ops = StaticTensorArrayOps(p, dtype, shape)\n",
    "    static_tensor_array_ops.register()\n",
    "\n",
    "    tensor_array = p.get_global_var_static(\"tensor_array\", dtype, shape)\n",
    "    tensor = p.get_tensor_ctor_static(\"tensor_constructor\", dtype, shape)\n",
    "    write = p.get_global_var_static(\"tensor_array_write\", dtype, shape)\n",
    "    gather = p.get_global_var_static(\"tensor_array_gather\", dtype, shape)\n",
    "    v = relay.var(\"v\")\n",
    "    indice = relay.var(\"indice\")\n",
    "    init_tensor_array = tensor_array(relay.const(3))\n",
    "    tensor_array1 = write(init_tensor_array, relay.const(0), tensor(v))\n",
    "    tensor_array2 = write(tensor_array1, relay.const(1), tensor(v))\n",
    "    tensor_array3 = write(tensor_array2, relay.const(2), tensor(v))\n",
    "    out = gather(tensor_array3, indice)\n",
    "    mod[\"main\"] = relay.Function([v, indice], out)\n",
    "    mod = transform.RemoveUnusedFunctions()(mod)\n",
    "    mod = transform.PartitionGraph()(mod)\n",
    "\n",
    "run(\"float32\", [2, 3])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 测试不绑定常量"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "False"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def get_net(prefix, data, out_channel):\n",
    "    weight = relay.var(prefix + \"weight\")\n",
    "    bn_gamma = relay.var(prefix + \"bn_gamma\")\n",
    "    bn_beta = relay.var(prefix + \"bn_beta\")\n",
    "    bn_mmean = relay.var(prefix + \"bn_mean\")\n",
    "    bn_mvar = relay.var(prefix + \"bn_var\")\n",
    "\n",
    "    layer = relay.nn.conv2d(\n",
    "        data=data, weight=weight, kernel_size=(3, 3), channels=out_channel, padding=(1, 1)\n",
    "    )\n",
    "    bn_output = relay.nn.batch_norm(layer, bn_gamma, bn_beta, bn_mmean, bn_mvar)\n",
    "    out = relay.nn.relu(bn_output[0])\n",
    "    return relay.Function(relay.analysis.free_vars(out), out)\n",
    "\n",
    "def get_partitoned_mod(mod, params, pattern_table, bind_constants):\n",
    "    mod[\"main\"] = bind_params_by_name(mod[\"main\"], params)\n",
    "    remove_bn_pass = tvm.transform.Sequential(\n",
    "        [\n",
    "            transform.InferType(),\n",
    "            transform.SimplifyInference(),\n",
    "            transform.FoldConstant(),\n",
    "            transform.FoldScaleAxis(),\n",
    "        ]\n",
    "    )\n",
    "    composite_partition = tvm.transform.Sequential(\n",
    "        [\n",
    "            remove_bn_pass,\n",
    "            transform.MergeComposite(pattern_table),\n",
    "            transform.AnnotateTarget(\"dnnl\"),\n",
    "            transform.PartitionGraph(bind_constants=bind_constants),\n",
    "        ]\n",
    "    )\n",
    "\n",
    "    with tvm.transform.PassContext(opt_level=3, disabled_pass=[\"AlterOpLayout\"]):\n",
    "        return composite_partition(mod)\n",
    "\n",
    "data = relay.var(\"data\", relay.TensorType((1, 3, 224, 224), \"float32\"))\n",
    "net = get_net(\"block_\", data, 8)\n",
    "mod, params = tvm.relay.testing.create_workload(net)\n",
    "\n",
    "mod = get_partitoned_mod(mod, params, get_pattern_table(\"dnnl\"), bind_constants=True)\n",
    "len(mod[\"main\"].body.args) == 1\n",
    "\n",
    "mod = get_partitoned_mod(mod, params, get_pattern_table(\"dnnl\"), bind_constants=False)\n",
    "len(mod[\"main\"].body.args) == 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ai",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
